#include "HuffmanCompressor.h"
#include <fstream>
#include <iostream>
#include <filesystem>

namespace fs = std::filesystem;

// 设置压缩对象的位置
void HuffmanCompressor::setTargetLocation(const std::string& location) {
    targetLocation = location;
}

// 设置压缩包的位置
void HuffmanCompressor::setHuffLocation(const std::string& location) {
    huffLocation = location;
}

// 设置解压后的位置
void HuffmanCompressor::setUnHuffLocation(const std::string& location) {
    unHuffLocation = location;
}

// 获取待压缩文件的字节频率
std::map<unsigned char, long long> HuffmanCompressor::getWeights(const std::string& inputFile) {
    std::ifstream fin(inputFile, std::ios::in | std::ios::binary);  // 使用二进制模式打开文件
    std::map<unsigned char, long long> weights;

    const size_t bufferSize = 8192;  // 8KB 缓冲区
    std::vector<char> buffer(bufferSize);  // 创建缓冲区
    while (fin.read(buffer.data(), buffer.size()) || fin.gcount() > 0) {  // 读取文件到缓冲区
        for (std::streamsize i = 0; i < fin.gcount(); ++i) {
            unsigned char c = buffer[i];
            weights[c]++;  // 更新字节频率
        }
    }

    fin.close();
    return weights;
}

// 将二进制字符串转换为单字节字符
unsigned char HuffmanCompressor::binaryStringToChar(const std::string& binary) {
    int sum = 0;
    for (size_t i = 0; i < binary.size(); ++i) {
        if (binary[i] == '1') {
            sum += (1 << (binary.size() - 1 - i));  // 计算对应的二进制值
        }
    }
    return static_cast<unsigned char>(sum);
}

// 将单字节字符转换为二进制字符串
std::string HuffmanCompressor::ucharToBinaryString(unsigned char value) {
    std::string result;
    for (int i = 7; i >= 0; --i) {
        result += (value & (1 << i)) ? '1' : '0';  // 从高位到低位逐位转换
    }
    return result;
}

// 遍历目标并收集目录和文件信息
void HuffmanCompressor::collectFolderInfo(const std::string& inputName, std::vector<std::string>& filePaths, std::vector<std::string>& folderPaths) {
    // 使用 std::filesystem::path 来处理路径，确保拼接正确
    fs::path folder(targetLocation + "\\" + inputName);

    // 如果是文件，直接将文件路径加入 filePaths
    if (std::filesystem::is_regular_file(folder)) {
        filePaths.push_back(inputName);  // 直接加入文件路径
        return;  // 如果是文件，就不再进行递归遍历
    }

    // 如果是文件夹，则遍历文件夹内容
    if (std::filesystem::is_directory(folder)) {
        // 遍历文件夹
        for (const auto& entry : fs::recursive_directory_iterator(folder)) {
            // 获取相对路径
            fs::path relativePath = fs::relative(entry.path(), folder);  // 获取相对路径

            if (entry.is_directory()) {
                folderPaths.push_back(relativePath.string());  // 存储相对路径
            }
            else if (entry.is_regular_file()) {
                filePaths.push_back(relativePath.string());  // 存储相对路径
            }
        }
    }
}

// 将目录结构信息写入文件头
void HuffmanCompressor::writeFolderMetadata(std::ofstream& fout, std::string folderName, const std::vector<std::string>& folderPaths, const std::vector<std::string>& filePaths) {
    fout << folderName << " ";  // 先写入根文件夹名
    fout << folderPaths.size() << " " << filePaths.size() << " ";  // 写入总子文件夹数和总文件数

    for (const auto& folder : folderPaths) {
        fout << folder << " ";  // 写入每个文件夹的路径
    }

    for (size_t i = 0; i < filePaths.size(); ++i) {
        fout << filePaths[i] << " " ;  // 写入每个文件的路径
    }
}

// 设置压缩密码
void HuffmanCompressor::setCompressPassword(std::ofstream& fout) {
    std::cout << "Set compress-password? (y/n): ";
    std::string choice;
    while(true) {
        std::cin >> choice;
        if (choice == "y") {
            fout << 1 << " ";
            std::cout << "Input the password: ";
            std::string password;
            std::cin >> password;
            fout << password << " ";
            break;
        }
        else if(choice == "n") {
            fout << 0 << " ";
            break;
        }
        else {
            std::cout << "Unknown command. Available commands: y, n." << std::endl;
            std::cout << "Input again: ";
        }
    }
}

// 处理文件覆盖
bool HuffmanCompressor::askForOverwrite(const std::string& fileName) {
    std::string choice;
    std::cout << "The file " << fileName << " already exists. Do you want to overwrite it? (y/n): ";
    while (true) {
        std::cin >> choice;
        if (choice == "y") {
            return true;
        }
        else if (choice == "n") {
            return false;
        }
        else {
            std::cout << "Unknown command. Available commands: y, n." << std::endl;
            std::cout << "Input again: ";
        }
    }
}

// 压缩函数
void HuffmanCompressor::compressFile(const std::string& inputFile, const std::string& compressedFile) {
    // 处理文件覆盖问题
    if (fs::exists(huffLocation + '\\' + compressedFile)) {
        bool reaction = askForOverwrite(compressedFile);
        if (!reaction) {
            std::cout << "Compression aborted. The file was not overwritten." << std::endl;
            return;
        }
        else {
            fs::remove(huffLocation + '\\' + compressedFile);  // 删除原文件
        }
    }
    
    // 二进制写，在指定位置下创建压缩文件
    std::ofstream fout(huffLocation+ '\\' + compressedFile, std::ios::out | std::ios::binary);
    
    // 设置压缩密码
    setCompressPassword(fout);

    // 定义缓冲区大小
    const size_t bufferSize = 8192;  // 8KB缓冲区

    // 收集目录信息
    std::vector<std::string> filePaths;
    std::vector<std::string> folderPaths;
    collectFolderInfo(inputFile, filePaths, folderPaths);

    // 放置一个标志位，是文件就为0，是目录就为1
    int flag;
    if (std::filesystem::is_regular_file(targetLocation + '\\' + inputFile))
        flag = 0;
    else
        flag = 1;
    fout << flag << " ";

    // 文件头写入目录和文件结构信息
    writeFolderMetadata(fout, inputFile, folderPaths, filePaths);

    // 遍历每个文件
    for (size_t i = 0; i < filePaths.size(); ++i) {
        const std::string& file = filePaths[i];
        std::string fileName;
        if (flag)
            fileName = targetLocation + '\\' + inputFile + "\\" + file;
        else
            fileName = targetLocation + '\\' + inputFile;

        // 检查空文件
        std::ifstream fin(fileName, std::ios::in | std::ios::binary);
        if (fin.peek() == std::ifstream::traits_type::eof()) {
            // 空文件处理
            fout << 0 << ' ';  // 空文件的频率表为空
            continue;
        }

        std::map<unsigned char, long long> weights = getWeights(fileName);

        HuffmanTree tree(weights);
        tree.generateTree();

        const auto& encodingMap = tree.getEncodingMap();
        std::string binary = "";
        std::vector<char> buffer(bufferSize);  // 创建缓冲区

        // 获取文件内容并进行哈夫曼编码
        while (fin.read(buffer.data(), bufferSize) || fin.gcount() > 0) {
            for (std::streamsize i = 0; i < fin.gcount(); ++i) {
                binary += encodingMap.at(static_cast<unsigned char>(buffer[i]));  // 根据字符获取哈夫曼编码
            }
        }
        fin.close();

        // 二进制字符串可能不是8的倍数，需要在末尾填充0
        int addZero = 8 - binary.size() % 8;
        for (int i = 0; i < addZero; ++i) binary += '0';

        // 记录该文件压缩后的字节数
        long long compressedSize = binary.size() / 8 + (addZero > 0 ? 1 : 0);
    
        // 写入文件头部分
        fout << weights.size() << " " << addZero << " ";
        for (const auto& entry : weights) {
            fout << entry.first << " " << entry.second << " ";
        }
        fout << compressedSize << " ";  // 写入当前文件大小

        // 创建一个缓冲区来批量写入字节
        std::vector<char> writeBuffer;
        // 将二进制数据转换为字节并加入到缓冲区
        for (size_t i = 0; i < binary.size(); i += 8) {
            unsigned char byte = binaryStringToChar(binary.substr(i, 8));  // 每8位二进制数据转换为一个字节
            writeBuffer.push_back(byte);  // 将字节添加到缓冲区
        }
        // 使用缓冲区写入文件
        fout.write(writeBuffer.data(), writeBuffer.size());  // 将整个缓冲区的数据写入文件
        // 如果有多个文件，则写入分隔符
        if (i < filePaths.size() - 1) fout << ' ';
    }

    fout.close();
    std::cout << "compress done" << std::endl;
}

// 解压缩函数
void HuffmanCompressor::decompressFile(const std::string& compressedFile) {
    std::ifstream fin(huffLocation+ '\\' + compressedFile, std::ios::in | std::ios::binary);

    //读取加密信息
    int isEncrypted;
    fin >> isEncrypted;
    if (isEncrypted) {
        std::string realPassword, inputPassword;
        fin >> realPassword;
        std::cout << "This compressed file has been encrypted, please input the password to decompress it: ";
        while (true){
            std::cin >> inputPassword;
            if (realPassword == inputPassword) {
                std::cout << "Correct password." << std::endl;
                break;
            }
            else {
                std::cout << "Wrong password." << std::endl;
                std::cout << "Input again: ";
            }
        }
    }

    // 读取标志位
    int flag;
    fin >> flag;

    // 读取文件头：根目录名、子文件夹数量、子文件数量、文件夹路径、文件路径及压缩大小
    std::string rootName;
    int folderCount, fileCount;
    fin >> rootName >> folderCount >> fileCount;

    std::vector<std::string> folderPaths(folderCount);
    std::vector<std::string> filePaths(fileCount);

    for (int i = 0; i < folderCount; ++i) {
        fin >> folderPaths[i];  // 读取子文件夹路径
    }

    for (int i = 0; i < fileCount; ++i) {
        fin >> filePaths[i];  // 读取文件路径
    }

    // 创建根目录
    fs::path rootPath = unHuffLocation + "\\" + rootName;
    // 处理文件覆盖问题
    if (fs::exists(rootPath)) {
        bool reaction = askForOverwrite(rootName);
        if (!reaction) {
            std::cout << "DeCompression aborted. The file was not overwritten." << std::endl;
            return;
        }
        else {
            if (fs::is_regular_file(rootPath)) fs::remove(rootPath);
            else if (fs::is_directory(rootPath)) fs::remove_all(rootPath);
        }
    }
    if (flag) {
        fs::create_directory(rootPath);  // 如果是目录，则创建
    }

    // 创建子文件夹
    for (const auto& folderPath : folderPaths) {
        fs::path folderDir = rootPath / folderPath;
        fs::create_directory(folderDir);  // 如果子文件夹不存在，则创建
    }

    // 解压每个文件
    for (int i = 0; i < fileCount; ++i) {
        const auto& file = filePaths[i];

        int size, addZero;
        fin >> size;
        // 如果是空文件，跳过解压
        if (size == 0) {
            fs::path filePath;
            if (flag)
                filePath = rootPath / file;
            else
                filePath = rootPath;
            std::ofstream foutFile(filePath, std::ios::out | std::ios::binary);  // 创建空文件
            foutFile.close();
            continue;
        }
        fin >> addZero;
        fin.get();  // 读出空格

        // 获取字节频率并构建哈夫曼树
        std::map<unsigned char, long long> weights;
        for (int j = 0; j < size; ++j) {
            unsigned char c;
            long long weight;
            fin.get((char&)c);  // 读取字节
            fin >> weight;
            fin.get();  // 跳过空格或换行
            weights[c] = weight;
        }

        HuffmanTree tree(weights);
        tree.generateTree();
        const auto& huffmanTree = tree.getTree();

        // 读取压缩的二进制数据
        std::string binary = "";
        long long compressedSize;  // 获取当前文件的压缩大小
        fin >> compressedSize;
        fin.get();
        // 使用一个缓冲区来批量读取数据
        const size_t bufferSize = 8 * 1024;  // 缓冲区大小为8kB
        std::vector<unsigned char> buffer(bufferSize);
        // 读取并处理文件内容
        long long bytesRead = 0;
        while (bytesRead < compressedSize - 1) {
            size_t toRead = std::min(bufferSize, static_cast<size_t>(compressedSize - 1 - bytesRead));  // 确保不会超过剩余的压缩大小
            // 读取数据到缓冲区
            fin.read(reinterpret_cast<char*>(buffer.data()), toRead);
            size_t actualRead = fin.gcount();  // 实际读取的字节数
            // 将每个字节转换为二进制字符串并追加到 binary 字符串中
            for (size_t i = 0; i < actualRead; ++i) {
                binary += ucharToBinaryString(buffer[i]);
            }
            bytesRead += actualRead;  // 更新已读取的字节数
        }

        // 去掉填充的零
        if (addZero > 0) {
            binary = binary.substr(0, binary.size() - addZero); 
        }

        // 创建文件并解压数据
        fs::path filePath;
        if (flag)
            filePath = rootPath / file;
        else
            filePath = rootPath;
        std::ofstream foutFile(filePath, std::ios::out | std::ios::binary);  // 打开每个文件
        if (!foutFile) {
            std::cerr << "unable to create file: " << filePath << std::endl;
            continue;
        }

        std::vector<char> buffer_write(bufferSize);  // 创建一个缓冲区
        size_t bufferIndex = 0;  // 缓冲区的写入索引

        int root = huffmanTree.size() - 1;
        int current = root;

        // 解码并恢复文件
        for (char bit : binary) {
            current = (bit == '0') ? huffmanTree[current].lc : huffmanTree[current].rc;

            // 匹配到一个字符，输出并重新开始匹配下一个字符
            if (huffmanTree[current].lc == 0 && huffmanTree[current].rc == 0) {
                buffer_write[bufferIndex++] = huffmanTree[current].ch;  // 将字符写入缓冲区

                // 如果缓冲区满了，就将缓冲区中的数据写入文件
                if (bufferIndex == bufferSize) {
                    foutFile.write(buffer_write.data(), bufferIndex);  // 批量写入
                    bufferIndex = 0;  // 重置缓冲区索引
                }

                current = root;  // 回到根节点，继续解码下一个字符
            }
        }

        // 写入剩余的字符（如果有的话）
        if (bufferIndex > 0) {
            foutFile.write(buffer_write.data(), bufferIndex);  // 写入剩余的数据
        }

        foutFile.close();  // 关闭每个解压的文件
    }

    fin.close();
    std::cout << "decompress done" << std::endl;
}
